import xml.etree.ElementTree as ET
import json
import zipfile
import os
import shutil
import uuid
import argparse
import re
import html
import math

def parse_html_to_ops(raw_html):
    if not raw_html:
        return None, None

    newline_placeholder = '___NEWLINE__'
    text_with_newlines = re.sub(r'<(br|/p)s*?/?>', newline_placeholder, raw_html, flags=re.IGNORECASE)
    cleanr = re.compile('<.*?>')
    plain_text_cleaned = html.unescape(re.sub(cleanr, '', text_with_newlines))
    
    lines = [line.strip() for line in plain_text_cleaned.split(newline_placeholder)]
    non_empty_lines = [line for line in lines if line]
    final_plain_text = '\n'.join(non_empty_lines).strip() or None

    ops = []
    try:
        xml_root = ET.fromstring(f"<body>{raw_html}</body>")
        
        def process_element(element, current_attrs):
            if element.text:
                text = html.unescape(element.text)
                if text:
                    ops.append({'insert': text, 'attributes': current_attrs.copy()})

            for child in element:
                child_attrs = current_attrs.copy()
                
                if child.tag.lower() == 'span' and 'style' in child.attrib:
                    style = child.attrib['style']
                    color_match = re.search(r'color:\s*([^;]+)', style)
                    if color_match:
                        child_attrs['color'] = color_match.group(1).strip().lower()

                if child.tag.lower() == 'b' or child.tag.lower() == 'strong':
                    child_attrs['bold'] = True
                
                process_element(child, child_attrs)

                if child.tail:
                    tail_text = html.unescape(child.tail)
                    if tail_text:
                        ops.append({'insert': tail_text, 'attributes': current_attrs.copy()})
        
        process_element(xml_root, {})

        final_ops = []
        for op in ops:
            parts = op['insert'].split(newline_placeholder)
            for i, part in enumerate(parts):
                if part:
                    final_ops.append({'insert': part, 'attributes': op['attributes']})
                if i < len(parts) - 1:
                    final_ops.append({'insert': '\n'})

        if final_ops:
            merged_ops = [final_ops[0]]
            for i in range(1, len(final_ops)):
                current_op = final_ops[i]
                last_op = merged_ops[-1]
                
                if (current_op.get('attributes') == last_op.get('attributes') and
                        current_op.get('insert') != '\n' and last_op.get('insert') != '\n'):
                    last_op['insert'] += current_op['insert']
                else:
                    merged_ops.append(current_op)
            final_ops = merged_ops

        if final_ops:
            if not final_ops[-1].get('insert', '').endswith('\n'):
                final_ops.append({'insert': '\n'})
            return final_plain_text, final_ops

    except ET.ParseError:
        return final_plain_text, None
        
    return final_plain_text, None



def get_node_properties(xml_node):
    data = {
        'text': None, 'note': None, 'hyperlink': None, 
        'image_id': None, 'image_title': None, 'rich_text_ops': None,
        'pos_x': 0.0, 'pos_y': 0.0
    }
    if xml_node is None:
        return data

    geo = xml_node.find('mxGeometry')
    if geo is not None:
        data['pos_x'] = float(geo.attrib.get('x', 0.0))
        data['pos_y'] = float(geo.attrib.get('y', 0.0))

    properties = xml_node.find('properties')
    if properties is not None:
        html_label_prop = properties.find(".//property[@key='com.thinkbuzan.gaia.entities.HTMLLabel']/HTMLLabel")
        if html_label_prop is not None and html_label_prop.text:
            data['text'], data['rich_text_ops'] = parse_html_to_ops(html_label_prop.text)
        
        if data.get('text') is None:
            text_prop = properties.find(".//property[@key='com.thinkbuzan.gaia.cell.text']")
            if text_prop is not None:
                text = text_prop.attrib.get('value')
                if text is None:
                    text = text_prop.text
                if text:
                    data['text'] = text.strip()

        notes_prop = properties.find(".//property[@key='com.thinkbuzan.gaia.cell.notes']/Notes")
        if notes_prop is not None and notes_prop.text:
            data['note'], _ = parse_html_to_ops(notes_prop.text)

        hyperlink_prop = properties.find(".//propertylist[@key='com.thinkbuzan.gaia.cell.hyperlinks']/property/hyperlink/link")
        if hyperlink_prop is not None and hyperlink_prop.text:
            data['hyperlink'] = hyperlink_prop.text.strip()

    style = xml_node.attrib.get('style')
    if style:
        match = re.search(r'image=([^;]+)', style)
        if match and xml_node.tag != 'floatingIdea':
            data['image_id'] = match.group(1)
            data['image_title'] = f"{match.group(1)}.png"
    
    image_tag = xml_node.find('image')
    if image_tag is not None:
        style = image_tag.attrib.get('style')
        if style:
            match = re.search(r'image=([^;]+)', style)
            if match:
                data['image_id'] = match.group(1)
                data['image_title'] = f"{match.group(1)}.png"

    return data

def build_zxm_node(properties, attachments, layout_pos=None):
    text = properties.get('text') or ''
    
    zxm_node = {
        'data': {
            'id': str(uuid.uuid4()),
            'text': text + '\n',
        },
        'children': []
    }

    if layout_pos is not None:
        zxm_node['data']['layout'] = layout_pos

    if properties.get('rich_text_ops'):
        zxm_node['data']['richText'] = {'ops': properties['rich_text_ops']}
    else:
        zxm_node['data']['richText'] = {'ops': [{'insert': text + '\n'}]}

    if properties.get('note'):
        zxm_node['data']['note'] = properties['note']
    if properties.get('hyperlink'):
        link = properties['hyperlink']
        if not re.match(r'^[a-zA-Z]+://', link):
            link = 'http://' + link
        zxm_node['data']['hyperlink'] = link
        zxm_node['data']['hyperlinkTitle'] = link
    
    image_id = properties.get('image_id')
    if image_id:
        new_image_name_with_ext = f"{image_id}.png"
        
        zxm_node['data']['resourceId'] = image_id
        zxm_node['data']['image'] = f'resources/{new_image_name_with_ext}'
        zxm_node['data']['imageTitle'] = properties.get('image_title', 'image')
        zxm_node['data']['imageSize'] = {'width': 131, 'height': 200}
        
        if image_id not in attachments:
            attachments.append(image_id)
        
        style_defaults = {
            "shape": "underline", "stroke": "", "font-size": 14,
            "font-family": "\u5fae\u8f6f\u96c5\u9ed1, \"Microsoft YaHei\"", "background": "transparent",
            "color": "#1A1A1A", "stroke-width": 3, "connect-type": "mind-curve",
            "connect-color": "#FF9754", "connect-width": 3
        }
        zxm_node['data'].update(style_defaults)

    return zxm_node

def imx_to_zxm(imx_path, zxm_path):
    temp_dir_imx = f'imx_temp_{uuid.uuid4()}'
    temp_dir_zxm = f'zxm_temp_{uuid.uuid4()}'
    os.makedirs(temp_dir_imx, exist_ok=True)
    os.makedirs(os.path.join(temp_dir_zxm, 'resources'), exist_ok=True)

    try:
        with zipfile.ZipFile(imx_path, 'r') as zip_ref:
            zip_ref.extractall(temp_dir_imx)

        imx_data_dir = os.path.join(temp_dir_imx, 'data')
        zxm_resources_dir = os.path.join(temp_dir_zxm, 'resources')
        if os.path.exists(imx_data_dir):
            for original_id in os.listdir(imx_data_dir):
                shutil.copy(os.path.join(imx_data_dir, original_id), os.path.join(zxm_resources_dir, f"{original_id}.png"))

        tree = ET.parse(os.path.join(temp_dir_imx, 'data.xml'))
        xml_root = tree.getroot()

        node_data_store = {node.attrib['id']: get_node_properties(node) for node in xml_root.findall('.//*[@id]')}

        central_idea_xml = xml_root.find('.//floatingIdea')
        if central_idea_xml is None:
            central_idea_xml = xml_root.find('.//branch')
        if central_idea_xml is None:
            raise ValueError("Cannot find central idea in data.xml")
        
        attachments = []
        root_props = node_data_store[central_idea_xml.attrib['id']]
        zxm_root = build_zxm_node(root_props, attachments)

        imx_id_to_zxm_node = {central_idea_xml.attrib['id']: zxm_root}

        queue = [(central_idea_xml, zxm_root)]

        processed_branches = set()

        while queue:
            parent_xml_node, parent_zxm_node = queue.pop(0)
            parent_imx_id = parent_xml_node.attrib['id']
            parent_pos = node_data_store[parent_imx_id]

            child_branches_xml = xml_root.findall(f'.//branch[@source="{parent_imx_id}"]')
            
            sorted_branches = []
            for branch_xml in child_branches_xml:
                if branch_xml.attrib['id'] in processed_branches:
                    continue
                processed_branches.add(branch_xml.attrib['id'])

                child_imx_id = branch_xml.attrib['target']
                target_node = xml_root.find(f'.//*[@id="{child_imx_id}"]')
                if target_node is not None:
                    target_props = node_data_store[child_imx_id]
                    
                    parent_x = parent_pos['pos_x']
                    parent_y = parent_pos['pos_y']
                    
                    child_x = target_props['pos_x']
                    child_y = target_props['pos_y']

                    dx = child_x - parent_x
                    dy = child_y - parent_y
                    angle = math.atan2(dy, dx)

                    sorted_branches.append((angle, branch_xml, target_node))
            
            sorted_branches.sort(key=lambda x: x[0])

            for angle, branch_xml, target_xml_node in sorted_branches:
                branch_id = branch_xml.attrib['id']
                child_imx_id = target_xml_node.attrib['id']
                
                branch_props = get_node_properties(branch_xml)
                target_node_props = node_data_store[child_imx_id]

                combined_props = target_node_props.copy()
                for key, value in branch_props.items():
                    if value is not None:
                        combined_props[key] = value
                
                image_on_edge = xml_root.find(f".//image[@parent='{branch_id}']")
                if image_on_edge is not None:
                    image_props = node_data_store.get(image_on_edge.attrib['id'], {})
                    if image_props.get('image_id'):
                        combined_props['image_id'] = image_props['image_id']
                        combined_props['image_title'] = image_props['image_title']

                layout_pos = {"side": "right" if angle > -math.pi/2 and angle < math.pi/2 else "left"}

                child_zxm_node = build_zxm_node(combined_props, attachments, layout_pos)
                parent_zxm_node['children'].append(child_zxm_node)
                
                imx_id_to_zxm_node[child_imx_id] = child_zxm_node
                queue.append((target_xml_node, child_zxm_node))

        relative_links = []
        imx_id_to_zxm_id = {imx_id: node['data']['id'] for imx_id, node in imx_id_to_zxm_node.items()}
        arrows = xml_root.findall('.//arrow[@source][@target]')
        for arrow in arrows:
            source_id = arrow.attrib['source']
            target_id = arrow.attrib['target']
            if source_id in imx_id_to_zxm_id and target_id in imx_id_to_zxm_id:
                link = {
                    "id": str(uuid.uuid4()), "text": "\u5173\u8054",
                    "start": {"nodeId": imx_id_to_zxm_id[source_id], "base": {"x":0, "y":0}, "draw": {"x":0, "y":0}, "control": {"x":0, "y":0}},
                    "end": {"nodeId": imx_id_to_zxm_id[target_id], "base": {"x":0, "y":0}, "draw": {"x":0, "y":0}, "control": {"x":0, "y":0}},
                    "style": {}, "richText": {"ops":[{"attributes":{},"insert":"\u5173\u8054"},{"insert":"\n","attributes":{}}]}, "isClean": True
                }
                relative_links.append(link)

        content_json = {
            'root': zxm_root, 'template': 'map', 'theme': 'classical18',
            'attachments': [], 'relativeLinks': relative_links
        }

        with open(os.path.join(temp_dir_zxm, 'content.json'), 'w', encoding='utf-8') as f:
            json.dump(content_json, f, ensure_ascii=False, indent=2)

        with open(os.path.join(temp_dir_zxm, 'metadata.json'), 'w', encoding='utf-8') as f:
            json.dump({'type': 'mindmap', 'template': 'map', 'creator': 'imx-converter'}, f)

        with open(os.path.join(temp_dir_zxm, 'manifest.json'), 'w', encoding='utf-8') as f:
            json.dump({}, f)

        with zipfile.ZipFile(zxm_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for root_dir, _, files in os.walk(temp_dir_zxm):
                for file in files:
                    zipf.write(os.path.join(root_dir, file), os.path.relpath(os.path.join(root_dir, file), temp_dir_zxm))

        print(f"Successfully converted {imx_path} to {zxm_path}")

    finally:
        if os.path.exists(temp_dir_imx):
            shutil.rmtree(temp_dir_imx)
        if os.path.exists(temp_dir_zxm):
            shutil.rmtree(temp_dir_zxm)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Convert iMindMap .imx files to .zxm files.')
    parser.add_argument('input', help='Input .imx file path')
    parser.add_argument('output', help='Output .zxm file path')
    args = parser.parse_args()
    imx_to_zxm(args.input, args.output)